package infra

import (
	"context"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

type transCtx struct{}

func FromTrans(ctx context.Context) (interface{}, bool) {
	v := ctx.Value(transCtx{})
	return v, v != nil
}

type noTransCtx struct {
}

func FromNoTrans(ctx context.Context) bool {
	v := ctx.Value(noTransCtx{})
	return v != nil && v.(bool)
}

type transLockCtx struct {
}

func FromTransLock(ctx context.Context) bool {
	v := ctx.Value(transLockCtx{})
	return v != nil && v.(bool)
}

func GetDB(ctx context.Context, defDB *gorm.DB) *gorm.DB {
	trans, ok := FromTrans(ctx)
	if ok && !FromNoTrans(ctx) {
		db, ok := trans.(*gorm.DB)
		if ok {
			if FromTransLock(ctx) {
				db = db.Clauses(clause.Locking{Strength: "UPDATE"})
			}
			return db
		}
	}

	return defDB
}

func NewTrans(ctx context.Context, trans interface{}) context.Context {
	return context.WithValue(ctx, transCtx{}, trans)
}
